import pygame
import random
import sys

# 初始化
pygame.init()
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇")
clock = pygame.time.Clock()

# 颜色定义
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

# 蛇初始位置
snake = [(100, 100), (90, 100), (80, 100)]
direction = "RIGHT"
food = (random.randint(0, WIDTH//10-1)*10, random.randint(0, HEIGHT//10-1)*10)

# 主循环
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        # 键盘控制
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and direction != "DOWN":
                direction = "UP"
            elif event.key == pygame.K_DOWN and direction != "UP":
                direction = "DOWN"
            elif event.key == pygame.K_LEFT and direction != "RIGHT":
                direction = "LEFT"
            elif event.key == pygame.K_RIGHT and direction != "LEFT":
                direction = "RIGHT"

    # 移动蛇头
    head_x, head_y = snake[0]
    if direction == "UP": head_y -= 10
    elif direction == "DOWN": head_y += 10
    elif direction == "LEFT": head_x -= 10
    elif direction == "RIGHT": head_x += 10
    new_head = (head_x, head_y)

    # 检测碰撞
    if (new_head in snake or
        new_head[0] < 0 or new_head[0] >= WIDTH or
        new_head[1] < 0 or new_head[1] >= HEIGHT):
        pygame.quit()
        sys.exit()

    # 添加新头部
    snake.insert(0, new_head)

    # 食物逻辑
    if snake[0] == food:
        food = (random.randint(0, WIDTH//10-1)*10,
                random.randint(0, HEIGHT//10-1)*10)
    else:
        snake.pop()  # 删除尾部

    # 渲染
    screen.fill(BLACK)
    for segment in snake:
        pygame.draw.rect(screen, GREEN, (segment[0], segment[1], 10, 10))
    pygame.draw.rect(screen, RED, (food[0], food[1], 10, 10))
    pygame.display.update()
    clock.tick(15)  # 控制帧率
